import torch
import torch.nn as nn

class TransformerModel(nn.Module):
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, nhead, num_encoder_layers,
                 num_decoder_layers, dim_feedforward, max_encoder_length, max_decoder_length, dropout=0.1):
        super(TransformerModel, self).__init__()

        # 定义源语言和目标语言的嵌入层
        self.src_embed = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embed = nn.Embedding(tgt_vocab_size, d_model)

        # 定义位置编码（这里为了简化省略了，实际应用中需要添加）

        # 定义编码器和解码器
        self.encoder = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout),
            num_layers=num_encoder_layers
        )

        self.decoder = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout),
            num_layers=num_decoder_layers
        )

        # 定义一个输出层，用于将解码器的输出映射到目标词汇表的大小
        self.output_layer = nn.Linear(d_model, tgt_vocab_size)

        # 初始化参数
        self.init_weights()

        # 这些参数用于生成注意力掩码（可选）
        self.max_encoder_length = max_encoder_length
        self.max_decoder_length = max_decoder_length

    def init_weights(self):
        """ 初始化模型参数 """
        initrange = 0.1
        self.src_embed.weight.data.uniform_(-initrange, initrange)
        self.tgt_embed.weight.data.uniform_(-initrange, initrange)
        self.output_layer.weight.data.uniform_(-initrange, initrange)
        self.output_layer.bias.data.zero_()

    def forward(self, src, tgt, src_mask=None, tgt_mask=None, src_key_padding_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None):
        r"""
        Args:
            src: Tensor, shape [seq_len, batch_size]
            tgt: Tensor, shape [tgt_seq_len, batch_size]
            src_mask: optional tensor, shape [seq_len, seq_len]
            tgt_mask: optional tensor, shape [tgt_seq_len, tgt_seq_len]
            src_key_padding_mask: Tensor, shape [batch_size, seq_len], optional
            tgt_key_padding_mask: Tensor, shape [batch_size, tgt_seq_len], optional
            memory_key_padding_mask: Tensor, shape [batch_size, src_seq_len], optional

        Returns:
            output: Tensor, shape [tgt_seq_len, batch_size, tgt_vocab_size]
        """
        # 将输入序列转换为嵌入表示
        src = self.src_embed(src) * math.sqrt(self.d_model)
        tgt = self.tgt_embed(tgt) * math.sqrt(self.d_model)

        # 如果提供了位置编码，则在此处添加

        # 通过编码器
        memory = self.encoder(src, src_mask=src_mask, src_key_padding_mask=src_key_padding_mask)

        # 通过解码器
        output = self.decoder(tgt, memory, tgt_mask=tgt_mask,
                             memory_mask=src_mask,
                             tgt_key_padding_mask=tgt_key_padding_mask,
                             memory_key_padding_mask=memory_key_padding_mask)

        # 将解码器的输出映射到目标词汇表的大小
        output = self.output_layer(output)

        return output

